#!/bin/env python
# -*- coding: utf-8 -*-

"""
File Name: region_client.py
Author: morre <morre@vip.qq.com>
Create Time: 2022/12/31 17:02:39
Brief:
"""

from __future__ import absolute_import
import logging
import socket
from struct import pack, unpack
from threading import Lock
from collections import deque, namedtuple
from io import BytesIO

from hbase_client.pb.RPC_pb2 import ConnectionHeader, RequestHeader, ResponseHeader
from hbase_client.pb.Client_pb2 import GetResponse, MutateResponse, ScanResponse, MultiResponse

# Used to encode and decode varints in a format protobuf expects.
from . import varint_util
from .hbase_exceptions import (
    RegionServerSocketTimeoutException, RegionServerException, PyBaseException,
    NoSuchColumnFamilyException, RegionMovedException, RegionOpeningException,
    RegionTooBusyException, DoNotRetryIOException)

logger = logging.getLogger("hbase_client." + __name__)
# logger.setLevel(logging.DEBUG)

# socket.setdefaulttimeout interfers with gevent.
# socket.setdefaulttimeout(2)

CurRegionActionResult = namedtuple(
    "CurRegionActionResult", ["result_or_exception", "exception"])

# We need to know how to interpret an incoming proto.Message. This maps
# the request_type to the response_type.
response_types = {
    "Get": GetResponse,
    "Mutate": MutateResponse,
    "Scan": ScanResponse,
    "Multi": MultiResponse
}


# This Client is created once per RegionServer. Handles all communication
# to and from this specific RegionServer.
class RegionClient():
    """Region服务客户端类"""
    # Variables are as follows:
    #   - Host: The hostname of the RegionServer
    #   - Port: The port of the RegionServer
    #   - Sock: An open connection to the RegionServer
    #   - call_id: A monotonically increasing int used as a sequence number
    #   for rpcs. This way we can match incoming responses with the rpc
    #   that made the request.

    def __init__(self, host, port, init_pool_size, socket_timeout):
        self.host = host  # string
        self.port = port  # int
        self.init_pool_size = init_pool_size
        self.socket_timeout = socket_timeout
        # We support connection pools so have lists of sockets and read/write
        # mutexes on them.
        self.sock_pool = deque()
        self.pool_lock = Lock()
        # We would like the region client to keep track of the regions that it
        # hosts. That way if we detect a Region server issue when touching
        # one region, we can close them all at the same time (saving us a
        # significant amount of meta lookups).
        self.regions = []
        self.call_id = 0
        self.call_lock = Lock()

    # Given an open socket, sends a ConnectionHeader over the wire to
    # initialize the connection.
    def send_hello(self, sock):
        ch = ConnectionHeader()
        ch.user_info.effective_user = b"hbase_client"
        ch.service_name = b"ClientService"
        serialized = ch.SerializeToString()
        # Message is serialized as follows -
        #   1. "HBas\x00\x50". Magic prefix that HBase requires.
        #   2. Little-endian uint32 indicating length of
        #      serialized ConnectionHeader
        #   3. Serialized ConnectionHeader
        message = b"HBas\x00\x50" + pack(">I", len(serialized)) + serialized
        sock.sendall(message)

    def _create_socket(self):
        s = None
        try:
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            logger.debug("connect %s:%d", self.host, self.port)
            s.connect((self.host, self.port))
            self.send_hello(s)
            s.settimeout(self.socket_timeout)
        except socket.timeout as ex:
            # 在hadoop上，如果send ok, recv失败。
            # 因此抛出异常重新更换调度节点。
            logger.exception(
                "create(%s:%d) socket timeout, except: %s",
                self.host, self.port, str(ex))
            raise RegionServerSocketTimeoutException(
                host=self.host, port=self.port, region_client=self) from ex
        except socket.error as ex:
            logger.exception(
                "create(%s:%d) socket error, except: %s",
                self.host, self.port, str(ex))
            raise RegionServerException(region_client=self) from ex
        return s

    # Creates the socket, initializes the connection.
    def open(self):
        is_open = True
        for _ in range(self.init_pool_size):
            s = self._create_socket()
            if not s:
                is_open = False
                break
            self.sock_pool.append(s)
        return is_open

    # morre: run means send + receive
    def run(self, rq):
        try:
            sock = None
            with self.pool_lock:
                if len(self.sock_pool) == 0:
                    sock = self._create_socket()
                else:
                    sock = self.sock_pool.popleft()

            rq = self.send(sock, rq)
            ret = self.receive(sock, rq)
        except PyBaseException as send_ex:
            if rq.type == "Multi":
                rpc = []
                tmp = CurRegionActionResult(
                    result_or_exception=[], exception=send_ex)
                rpc.append(tmp)
                return rpc
            raise send_ex

        with self.pool_lock:
            self.sock_pool.append(sock)
        return ret

    # Sends an RPC over the wire then calls _receive_rpc and returns the
    # response RPC.
    #
    # The raw bytes we send over the wire are composed (in order) -
    #
    #   1. little-endian uint32 representing the total-length of
    #      the following message.
    #   2. A single byte representing the length of
    #      the serialized RequestHeader.
    #   3. The serialized RequestHeader.
    #   4. A varint representing the length of the serialized RPC.
    #   5. The serialized RPC.
    #
    def send(self, sock, rq):
        # with self.call_lock:
        #    current_call_id = self.call_id
        #    self.call_id += 1
        serialized_rpc = rq.pb.SerializeToString()
        header = RequestHeader()
        #header.call_id = current_call_id
        header.method_name = rq.type
        header.request_param = True
        serialized_header = header.SerializeToString()
        # Consult the DESIGN.md for an explanation as to how Send/Receive
        # messages are composed.
        rpc_length_bytes = varint_util.to_varint(len(serialized_rpc))
        total_length = 4 + 1 + \
            len(serialized_header) + \
            len(rpc_length_bytes) + len(serialized_rpc)
        # Total length doesn"t include the initial 4 bytes (for the
        # total_length uint32)
        to_send = pack(b">IB", total_length - 4, len(serialized_header))
        to_send += serialized_header
        to_send += rpc_length_bytes
        to_send += serialized_rpc

        try:
            logger.debug(
                "Sending %s RPC to %s:%d", rq.type, self.host, self.port)
            sock.sendall(to_send)
        except socket.error as exc:
            logger.exception(
                "send region server(%s:%d) exception.", self.host, self.port)
            raise RegionServerException(region_client=self) from exc
        return rq

    # Called after sending an RPC, listens for the response and builds the
    # correct pbResponse object.
    #
    # The raw bytes we receive are composed (in order) -
    #
    #   1. little-endian uint32 representing the total-length of
    #      the following message.
    #   2. A varint representing the length of the serialized ResponseHeader.
    #   3. The serialized ResponseHeader.
    #   4. A varint representing the length of the serialized ResponseMessage.
    #   5. The ResponseMessage.
    #
    def receive(self, sock, rq, data=None):
        # TODO: 不应该在这里根据rq.type == "Multi"来返回
        # 应该在Multi的link_exception中处理。
        rpc = []
        recv_ex = None

        # If the field data is populated that means we should process from that
        # instead of the socket.
        full_data = data
        if data is None:
            try:
                msg_length = self._recv_n(sock, 4)
                if msg_length is None:
                    raise socket.error()
                msg_length = unpack(b">I", msg_length)[0]
                # The message is then going to be however many bytes the first
                # four bytes specified. We don"t want to overread or underread
                # as that will cause havoc.
                full_data = self._recv_n(
                    sock, msg_length)
            except socket.timeout:
                # 在hadoop上，如果send ok, recv失败。
                # 需要抛出异常重新更换调度节点。
                logger.exception(
                    "receive socket(%s:%d) timeout.", self.host, self.port)
                recv_ex = RegionServerSocketTimeoutException(
                    host=self.host, port=self.port, region_client=self)
            except socket.error:
                logger.exception(
                    "receive socket(%s:%d) error.", self.host, self.port)
                recv_ex = RegionServerException(region_client=self)

            if recv_ex:
                if rq.type == "Multi":
                    tmp = CurRegionActionResult(result_or_exception=[],
                                                exception=recv_ex)
                    rpc.append(tmp)
                    return rpc
                else:
                    raise recv_ex

        # Pass in the full data as well as your current position to the
        # decoder. It"ll then return two variables:
        #       - next_pos: The number of bytes of data specified by the varint
        #       - pos: The starting location of the data to read.
        next_pos, pos = varint_util.decodeVarint(full_data, 0)
        header = ResponseHeader()
        header.ParseFromString(full_data[pos: pos + next_pos])
        pos += next_pos

        if header.exception.exception_class_name != "":
            exception_class = header.exception.exception_class_name
            logger.warning("get java exception: %s", exception_class)
            raise self.exception_class2exception(
                exception_class, header.exception.stack_trace)
        next_pos, pos = varint_util.decodeVarint(full_data, pos)
        rpc = None
        resp = response_types[rq.type]()
        resp.ParseFromString(full_data[pos: pos + next_pos])

        if rq.type == "Multi":
            rpc = []
            for region_action_result in resp.regionActionResult:
                local_exception, result_or_exception = None, None
                if region_action_result.HasField("exception"):
                    logger.warning(
                        "region_action_result get exception:%s errmsg:%s",
                        region_action_result.exception.name,
                        region_action_result.exception.value)
                    local_exception = self.exception_class2exception(
                        region_action_result.exception.name,
                        region_action_result.exception.value)

                result_or_exception = []
                for result in region_action_result.resultOrException:
                    action_result = {
                        "index": result.index,
                        "result": None,
                        "exception": None,
                        "errmsg": None
                    }
                    if result.HasField("result"):
                        action_result["result"] = result.result
                    if result.HasField("exception"):
                        action_result["errmsg"] = result.exception.value
                        exception = result.exception.name
                        action_result["exception"] = \
                            self.exception_class2exception(
                                exception, result.exception.value)
                    result_or_exception.append(action_result)
                tmp = CurRegionActionResult(
                    result_or_exception=result_or_exception,
                    exception=local_exception)
                rpc.append(tmp)
        else:
            rpc = resp
        # The rpc is fully built!
        return rpc

    def exception_class2exception(self, exception_class, exc_info=None):
        if exception_class in [
            "org.apache.hadoop.hbase.regionserver.NoSuchColumnFamilyException",
            "java.io.IOException"]:
            return NoSuchColumnFamilyException()
        elif exception_class == \
            "org.apache.hadoop.hbase.exceptions.RegionMovedException":
            return RegionMovedException()
        elif exception_class == \
            "org.apache.hadoop.hbase.NotServingRegionException":
            # return NotServingRegionException()
            return RegionServerException(region_client=self)
        elif exception_class == \
            "org.apache.hadoop.hbase.regionserver.WrongRegionException":
            return RegionServerException(region_client=self)
        elif exception_class == \
            "org.apache.hadoop.hbase.regionserver.RegionServerStoppedException":
            return RegionServerException(region_client=self)
        elif exception_class == \
            "org.apache.hadoop.hbase.exceptions.RegionOpeningException":
            return RegionOpeningException()
        elif exception_class == \
            "org.apache.hadoop.hbase.RegionTooBusyException":
            return RegionTooBusyException()
        elif exception_class == \
            "org.apache.hadoop.hbase.ipc.ServerNotRunningYetException":
            return RegionServerException(region_client=self)
        elif exception_class == \
            "org.apache.hadoop.hbase.DoNotRetryIOException":
            return DoNotRetryIOException()
        elif exception_class == \
            "org.apache.hadoop.hbase.regionserver.RegionServerAbortedException":
            return RegionServerException(region_client=self)
        else:
            logging.error("Unknown exception: %s", exception_class)
            return PyBaseException(
                exception_class + f". Remote traceback:\n{exc_info}")

    # Receive an RPC with incorrect call_id?
    #       1. Acquire lock
    #       2. Place raw data into missed_rpcs with key call_id
    #       3. Notify all other threads to wake up (nothing will happen until
    #          you release the lock)
    #       4. WHILE: Your call_id is not in the dictionary
    #               4.5  Call wait() on the conditional and get comfy.
    #       5. Pop your data out
    #       6. Release the lock
    # def receive_other_call_id(self, call_id, req, msg_id, data):
    #    with self.missed_rpcs_lock:
    #        logger.debug(
    #            "Received invalid RPC ID. Got: %s, Expected: %s.",
    #            msg_id, call_id)
    #        self.missed_rpcs[msg_id] = data
    #        self.missed_rpcs_condition.notifyAll()
    #        while call_id not in self.missed_rpcs:
    #            if self.shutting_down:
    #                raise RegionServerException(region_client=self)
    #            self.missed_rpcs_condition.wait()
    #        new_data = self.missed_rpcs.pop(call_id)
    #        logger.debug("Another thread found my RPC! RPC ID: %s", call_id)
    #    return self.receive(call_id, req, data=new_data)

    # Receives exactly n bytes from the socket. Will block until n bytes are
    # received. If a socket is closed (RegionServer died) then raise an
    # exception that goes all the way back to the main client
    def _recv_n(self, sock, n):
        partial_str = BytesIO()
        partial_len = 0
        while partial_len < n:
            packet = sock.recv(n - partial_len)
            if not packet:
                raise socket.error()
            partial_len += len(packet)
            partial_str.write(packet)
        return partial_str.getvalue()

    # Do any work to close open file descriptors, etc.
    def close(self):
        for sock in self.sock_pool:
            sock.close()
        # We could still have greenlets waiting in the bad_call_id pools! Wake
        # them up so they can fail to error handling as well.
        # self.missed_rpcs_condition.acquire()
        # self.missed_rpcs_condition.notifyAll()
        # self.missed_rpcs_condition.release()

    def __repr__(self):
        return str({
            "host": self.host,
            "port": self.port,
        })
